#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <errno.h>
#include <glib.h>
#include <libnbd.h>
#include "utils.h"

// 守护进程的全局上下文结构
typedef struct {
    GHashTable* lba_map;         // 内存中的 LBA -> Hash 映射表
    const char* cache_dir;       // 本地块缓存目录
    const char* source_blob_dir; // M1产物的 blobs 目录 (模拟远程仓库)
    uint64_t disk_size;          // 虚拟磁盘总大小
} daemon_context_t;


/**
 * @brief 从 M1 生成的 lba_map.bin 文件加载映射关系到内存中的 GHashTable
 * @param map_file_path lba_map.bin 的路径
 * @return 成功则返回 GHashTable 指针, 失败返回 NULL
 */
static GHashTable* load_lba_map_from_file(const char* map_file_path) {
    LOG_INFO("开始加载 LBA Map 文件: %s", map_file_path);
    FILE* fp = fopen(map_file_path, "rb");
    if (!fp) {
        LOG_ERROR("无法打开 LBA Map 文件");
        return NULL;
    }

    GHashTable* lba_map = g_hash_table_new_full(g_int64_hash, g_int64_equal, g_free, g_free);
    if (!lba_map) {
        LOG_ERROR("创建 GHashTable 失败");
        fclose(fp);
        return NULL;
    }

    lba_map_entry_t entry;
    size_t entries_loaded = 0;
    while (fread(&entry, sizeof(lba_map_entry_t), 1, fp) == 1) {
        long long* lba_key = g_new(long long, 1);
        *lba_key = entry.lba;
        char* hash_value = g_strdup(entry.hash);
        g_hash_table_insert(lba_map, lba_key, hash_value);
        entries_loaded++;
    }

    fclose(fp);
    LOG_INFO("LBA Map 加载完成, 共加载 %zu 个条目", entries_loaded);
    return lba_map;
}

// ==============================================================================
//  核心逻辑 handle_pread 函数仍然完全不变
// ==============================================================================
static int handle_pread(void *user_data, void *buf, uint32_t count, uint64_t offset, uint32_t flags, int *err) {
    daemon_context_t *ctx = (daemon_context_t*)user_data;
    
    if (offset % CHUNK_SIZE != 0 || count != CHUNK_SIZE) {
        LOG_ERROR("收到非对齐的 NBD 请求! offset=%lu, count=%u", offset, count);
        *err = EINVAL;
        return -1;
    }
    
    long long lba = offset / CHUNK_SIZE;

    const char* hash = g_hash_table_lookup(ctx->lba_map, &lba);
    if (!hash) {
        LOG_INFO("LBA %lld 是一个稀疏块, 返回零数据", lba);
        memset(buf, 0, count);
        return 0;
    }

    char cache_path[1024];
    snprintf(cache_path, sizeof(cache_path), "%s/%s/%s", ctx->cache_dir, CHUNK_HASH_NAME, hash);

    FILE* fp = fopen(cache_path, "rb");
    if (fp) {
        LOG_INFO("缓存命中: LBA %lld -> Hash %s", lba, hash);
        if (fread(buf, 1, count, fp) != count) {
            LOG_ERROR("从缓存文件 %s 读取数据失败", cache_path);
            fclose(fp);
            *err = EIO;
            return -1;
        }
        fclose(fp);
        return 0;
    }

    LOG_INFO("缓存未命中: LBA %lld -> Hash %s, 准备从源拉取", lba, hash);
    
    char source_path[1024];
    snprintf(source_path, sizeof(source_path), "%s/%s/%s", ctx->source_blob_dir, CHUNK_HASH_NAME, hash);
    
    fp = fopen(source_path, "rb");
    if (!fp) {
        LOG_ERROR("数据块在源目录中也找不到! LBA: %lld, Hash: %s, 路径: %s", lba, hash, source_path);
        *err = EIO;
        return -1;
    }

    if (fread(buf, 1, count, fp) != count) {
        LOG_ERROR("从源文件 %s 读取数据失败", source_path);
        fclose(fp);
        *err = EIO;
        return -1;
    }
    fclose(fp);

    FILE *cache_fp = fopen(cache_path, "wb");
    if (!cache_fp) {
        LOG_ERROR("无法创建缓存文件: %s", cache_path);
        return 0;
    }
    
    if (fwrite(buf, 1, count, cache_fp) != count) {
        LOG_ERROR("写入缓存文件 %s 失败", cache_path);
    }
    fclose(cache_fp);
    LOG_INFO("已将数据块 %s 写入本地缓存", hash);

    return 0;
}

void print_usage(const char* prog) {
    fprintf(stderr, "用法: sudo %s -m <map_file> -b <source_blobs_dir> -c <cache_dir> -d <nbd_device> -s <disk_size_mb>\n", prog);
    fprintf(stderr, "  -m: M1产物的 lba_map.bin 文件路径 (必需)\n");
    fprintf(stderr, "  -b: M1产物的 blobs 目录路径 (必需, 模拟远程仓库)\n");
    fprintf(stderr, "  -c: 用于存放数据块的本地缓存目录 (必需, 会自动创建)\n");
    fprintf(stderr, "  -d: 要使用的 NBD 设备, 如 /dev/nbd0 (必需)\n");
    fprintf(stderr, "  -s: 虚拟磁盘大小, 单位MB (必需, 必须与M1构建时一致)\n");
}

int main(int argc, char *argv[]) {
    char *map_file = NULL;
    char *source_blobs = NULL;
    char *cache_dir = NULL;
    char *nbd_device = NULL;
    long disk_size_mb = 0;
    int opt;

    while ((opt = getopt(argc, argv, "m:b:c:d:s:")) != -1) {
        switch (opt) {
        case 'm': map_file = optarg; break;
        case 'b': source_blobs = optarg; break;
        case 'c': cache_dir = optarg; break;
        case 'd': nbd_device = optarg; break;
        case 's': disk_size_mb = atol(optarg); break;
        default:
            print_usage(argv[0]);
            return 1;
        }
    }

    if (!map_file || !source_blobs || !cache_dir || !nbd_device || disk_size_mb <= 0) {
        print_usage(argv[0]);
        return 1;
    }
    
    daemon_context_t context = {0};
    context.cache_dir = cache_dir;
    context.source_blob_dir = source_blobs;
    context.disk_size = (uint64_t)disk_size_mb * 1024 * 1024;

    char full_cache_path[1024];
    snprintf(full_cache_path, sizeof(full_cache_path), "%s/%s", cache_dir, CHUNK_HASH_NAME);
    char cmd[2048];
    snprintf(cmd, sizeof(cmd), "mkdir -p %s", full_cache_path);
    system(cmd);

    context.lba_map = load_lba_map_from_file(map_file);
    if (!context.lba_map) {
        LOG_ERROR("初始化 LBA Map 失败");
        return 1;
    }

    // ==============================================================================
    //  关键 API 修改点: 使用最古老、最兼容的 libnbd API
    // ==============================================================================
    struct nbd_handle *nbd = nbd_create();
    if (nbd == NULL) {
        LOG_ERROR("无法创建 NBD 句柄: %s", nbd_get_error());
        return 1;
    }
    
    // API 修改 1: 使用 nbd_set_size()。这个函数在旧版本中是存在的。
    // 如果这个函数也不存在，那说明 libnbd 版本实在太低，可能需要升级。
    nbd_set_size(nbd, context.disk_size);
    LOG_INFO("设置虚拟磁盘大小为 %ld MB", disk_size_mb);
    
    // API 修改 2: 使用 nbd_set_pread()，直接传递函数指针和上下文。
    nbd_set_pread(nbd, handle_pread, &context);
    LOG_INFO("已直接注册 NBD 读取 (pread) 回调函数");
    
    // API 修改 3: 使用 nbd_set_read_only(nbd, 1) 明确告知设备是只读的。
    nbd_set_read_only(nbd, 1);

    // API 修改 4: 使用 nbd_connect_command()。这是一个非常稳定和兼容的函数，
    // 它会在后台 fork 和 exec 一个 nbd-client 进程来连接内核。
    char *nbd_cmd[] = {
        "nbd-client",
        "-d", // -d 表示在连接后分离, 让我们的守护进程继续运行
        (char *)nbd_device,
        NULL
    };

    LOG_INFO("准备通过 nbd-client 命令连接到内核设备 %s...", nbd_device);
    if (nbd_connect_command(nbd, nbd_cmd) == -1) {
        LOG_ERROR("通过命令连接到设备 %s 失败: %s", nbd_device, nbd_get_error());
        nbd_close(nbd);
        return 1;
    }

    // API 修改 5: 使用 nbd_poll() 自己构建事件循环。这是最底层的循环方式。
    LOG_INFO("连接成功! ocim-nbd-daemon 正在运行... 按 Ctrl+C 退出。");
    while (1) {
        int r = nbd_poll(nbd, -1); // -1 表示无限期等待事件
        if (r == -1) {
            LOG_ERROR("NBD 运行循环出错: %s", nbd_get_error());
            break;
        }
    }

    // 清理资源
    LOG_INFO("正在关闭守护进程...");
    nbd_close(nbd);
    g_hash_table_destroy(context.lba_map);

    return 0;
}